package com.gxc.heap;

import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

/**
 * 数据流的中位数
 * 中位数是有序整数列表中的中间值。如果列表的大小是偶数，则没有中间值，中位数是两个中间值的平均值。
 *
 * 例如 arr = [2,3,4] 的中位数是 3 。
 * 例如 arr = [2,3] 的中位数是 (2 + 3) / 2 = 2.5 。
 * 实现 MedianFinder 类:
 *
 * MedianFinder() 初始化 MedianFinder 对象。
 *
 * void addNum(int num) 将数据流中的整数 num 添加到数据结构中。
 *
 * double findMedian() 返回到目前为止所有元素的中位数。与实际答案相差 10-5 以内的答案将被接受
 */
public class MedianFinder {

    private PriorityQueue<Integer> maxHeap;//存储小元素
    private PriorityQueue<Integer> minHeap;//存储大元素
    private Map<Integer, Integer> delayed; // 记录数字的出现次数
    private int maxHeapSize;
    private int minHeapSize;

    public MedianFinder() {
        maxHeap = new PriorityQueue<>((a, b) -> b-a);
        minHeap = new PriorityQueue<>();
        delayed = new HashMap<>();
        maxHeapSize = 0;
        minHeapSize = 0;
    }

    public void addNum(int num) {
        if (maxHeap.isEmpty() || num <= maxHeap.peek()) {
            maxHeap.offer(num);
            maxHeapSize++;
        } else {
            minHeap.offer(num);
            minHeapSize++;
        }
        balance();
    }

    private void balance() {
        if (maxHeapSize>minHeapSize+1) {
            minHeap.offer(maxHeap.poll());
            minHeapSize++;
            maxHeapSize--;
            pruneMax(); // 可能需要清理maxHeap的新堆顶
        } else if (minHeapSize > maxHeapSize) {
            maxHeap.offer(minHeap.poll());
            maxHeapSize++;
            minHeapSize--;
            pruneMin(); // 可能需要清理minHeap的新堆顶
        }
    }

    private void pruneMin() {
        while (!minHeap.isEmpty()) {
            int num = minHeap.peek();
            if (delayed.containsKey(num)) {
                delayed.put(num, delayed.get(num)-1);
                if (delayed.get(num) == 0) {
                    delayed.remove(num);
                }
                minHeap.poll();
            } else {
                break;
            }
        }
    }

    private void pruneMax() {
        while (!maxHeap.isEmpty()) {
            int num = maxHeap.peek();
            if (delayed.containsKey(num)) {
                delayed.put(num, delayed.get(num)-1);
                if (delayed.get(num) == 0) {
                    delayed.remove(num);
                }
                maxHeap.poll();
            } else {
                break;
            }
        }
    }

    public double findMedian() {
        if (minHeap.size() == maxHeap.size()) {
            return (minHeap.peek() + maxHeap.peek())/2.0;
        } else {
            return maxHeap.peek();
        }
    }

    public void remove(int num) {
        // 更新哈希表
        delayed.put(num, delayed.getOrDefault(num, 0) + 1);
        if (num <= maxHeap.peek()) {
            maxHeapSize--;
            if (num == maxHeap.peek()) {
                pruneMax();
            }
        } else  {
            minHeapSize--;
            if (num == minHeap.peek()) {
                pruneMin();
            }
        }
        balance();
    }
}
